from __future__ import print_function  # 导入 print_function 以确保兼容 Python 2 和 3
import os  # 用于文件路径操作
import cv2  # OpenCV库，用于读取和处理图像
import argparse  # 用于从命令行解析参数
import numpy as np  # 数组处理库
import torch  # PyTorch库
import torch.nn as nn  # PyTorch中的神经网络模块
import torch.nn.functional as F  # 包含激活函数等
import torch.optim as optim  # PyTorch中的优化器
from torch.utils.data import Dataset, DataLoader  # 数据集和数据加载器
from torchvision import transforms  # 用于图像预处理的工具

# 定义训练设置参数
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')  # 定义命令行参数解析器
parser.add_argument('--batch-size', type=int, default=64, metavar='N',  # 批处理大小参数
                    help='input batch size for training (default: 64)')  # 训练时每次处理的数据批次大小
parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',  # 测试时每次处理的批次大小
                    help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs', type=int, default=10, metavar='N',  # 训练周期数
                    help='number of epochs to train (default: 10)')  # 训练数据的遍历次数
parser.add_argument('--lr', type=float, default=0.01, metavar='LR',  # 学习率参数
                    help='learning rate (default: 0.01)')  # 控制模型更新的步幅
parser.add_argument('--momentum', type=float, default=0.5, metavar='M',  # 动量参数
                    help='SGD momentum (default: 0.5)')  # 在优化时缓冲梯度的变化趋势，提升训练稳定性
parser.add_argument('--no-cuda', action='store_true', default=False,  # 是否禁用 CUDA
                    help='disables CUDA training')  # 禁用 GPU（CUDA）训练的开关
parser.add_argument('--seed', type=int, default=1, metavar='S',  # 随机数种子
                    help='random seed (default: 1)')  # 为了确保实验结果的可重复性
parser.add_argument('--log-interval', type=int, default=2, metavar='N',  # 日志间隔
                    help='how many batches to wait before logging training status')  # 多少个批次记录一次训练状态
args = parser.parse_args(args=[])  # 解析参数
args.cuda = not args.no_cuda and torch.cuda.is_available()  # 根据是否有CUDA和用户选择来确定是否使用GPU

torch.manual_seed(args.seed)  # 设置随机种子，确保每次运行结果一致
if args.cuda:
    torch.cuda.manual_seed(args.seed)  # 如果使用CUDA，也为GPU设置随机种子

# 定义数据集路径和保存路径
data_train_path = r'E:\CV\dataset\dataset\train'  # 训练数据路径
data_test_path = r'E:\CV\dataset\dataset\test'  # 测试数据路径
train_x_path = r'E:\CV\dataset\dataset\train_x_data.npy'  # 训练集图片的保存路径
train_y_path = r'E:\CV\dataset\dataset\train_y_data.npy'  # 训练集标签的保存路径
test_x_path = r'E:\CV\dataset\dataset\test_x_data.npy'  # 测试集图片的保存路径
test_y_path = r'E:\CV\dataset\dataset\test_y_data.npy'  # 测试集标签的保存路径

# 定义数据生成函数，用于读取图片并生成训练和测试数据
def generateds(path):
    x, y_ = [], []  # 用于存放图片和标签
    all_files = os.listdir(path)  # 获取路径下的所有文件夹
    if len(all_files) == 0:
        print(f"No directories found in {path}")
        return x, y_

    # 遍历每个文件夹
    for each_file in all_files:
        print('each_file')
        file_path = os.path.join(path, each_file)
        if not os.path.isdir(file_path):
            print(f"{file_path} is not a valid directory, skipping")
            continue

        print(f"Processing directory: {file_path}")
        all_image = os.listdir(file_path)  # 获取每个文件夹中的所有图像
        if len(all_image) == 0:
            print(f"No images found in directory {file_path}")
            continue

        for image in all_image:  # 遍历每张图片
            img_path = os.path.join(file_path, image)
            print(f"Processing image: {img_path}")
            img = cv2.imread(img_path)  # 使用 OpenCV 读取图像
            if img is None:
                print(f"Image {img_path} could not be read")
                continue

            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 转换为灰度图
            x.append(img)  # 将图像添加到数据列表中

            # 标签定义，文件夹名决定类别
            if os.path.basename(each_file) == 'A':
                y_.append([0])
            elif os.path.basename(each_file) == 'B':
                y_.append([1])
            elif os.path.basename(each_file) == 'C':
                y_.append([2])
            elif os.path.basename(each_file) == 'D':
                y_.append([3])

    print(f"Total images loaded: {len(x)}")  # 打印加载的图像总数
    return np.array(x), np.array(y_).astype(np.float32)  # 将数据转换为NumPy数组并返回

# 定义数据预处理操作
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为Tensor
    transforms.Normalize((0.1307,), (0.3081,))  # 归一化操作
])

# 自定义数据集类
class MyDataset(Dataset):
    def __init__(self, data_npy, target_npy):
        self.data = data_npy  # 图像数据
        self.label = target_npy  # 标签数据
        self.transforms = transform  # 使用的预处理方法
        print(f"Dataset loaded: {len(self.data)} samples")  # 打印加载的样本数

    def __getitem__(self, index):
        hdct = self.data[index]  # 获取图像
        ldct = self.label[index]  # 获取标签
        ldct = torch.tensor(ldct)  # 将标签转换为Tensor
        hdct = self.transforms(hdct)  # 对图像进行预处理
        return hdct, ldct  # 返回图像和标签

    def __len__(self):
        return len(self.data)  # 返回数据集的大小

# 定义卷积神经网络结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 定义两个卷积层和一个全连接层
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)  # 输入通道1，输出通道10，卷积核大小5x5
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)  # 输入通道10，输出通道20，卷积核大小5x5
        self.conv2_drop = nn.Dropout2d()  # Dropout层用于防止过拟合
        self.fc1 = nn.Linear(500, 50)  # 全连接层，输入500个神经元，输出50个
        self.fc2 = nn.Linear(50, 4)  # 输出层，4个类别对应4个输出

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))  # 经过卷积、ReLU激活和最大池化
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))  # 第二个卷积层+池化+Dropout
        x = x.view(-1, 500)  # 展平图像数据为向量
        x = F.relu(self.fc1(x))  # 通过全连接层和ReLU激活
        x = F.dropout(x, training=self.training)  # Dropout以防止过拟合
        x = F.relu(self.fc2(x))  # 通过第二个全连接层
        return F.log_softmax(x, dim=1)  # 使用LogSoftmax作为输出

# 训练函数
def train(epoch, model, train_loader, optimizer):
    model.train()  # 设置模型为训练模式
    for batch_idx, (data, target) in enumerate(train_loader):  # 遍历训练数据
        if args.cuda:
            data, target = data.cuda(), target.cuda()  # 将数据和标签移动到GPU
        optimizer.zero_grad()  # 梯度置零
        output = model(data)  # 前向传播
        target = target.type(torch.LongTensor)  # 将标签转换为Long类型
        target = target.cuda() if args.cuda else target  # 将标签移至GPU
        target = target.view(-1)  # 展平标签
        loss = F.nll_loss(output, target)  # 计算负对数似然损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新模型参数
        if batch_idx % args.log_interval == 0:  # 每隔log_interval批次输出一次日志
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

# 测试函数
def test(model, ztest_loader):
    model.eval()  # 设置模型为评估模式
    test_loss = 0
    correct = 0
    with torch.no_grad():  # 禁用梯度计算
        for data, target in test_loader:  # 遍历测试数据
            if args.cuda:
                data, target = data.cuda(), target.cuda()  # 将数据移至GPU
            output = model(data)  # 前向传播
            target = target.type(torch.LongTensor)  # 转换标签为Long类型
            target = target.cuda() if args.cuda else target  # 将标签移至GPU
            target = target.view(-1)  # 展平标签
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # 计算损失
            pred = output.argmax(dim=1, keepdim=True)  # 选择概率最大的类别
            correct += pred.eq(target.view_as(pred)).sum().item()  # 统计正确预测的数量

    test_loss /= len(test_loader.dataset)  # 计算平均损失
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))  # 打印测试损失和准确率

# 主入口，确保在Windows系统下正确使用多进程
if __name__ == "__main__":
    # 如果预处理后的数据文件存在，则加载数据
    if os.path.exists(train_x_path) and os.path.exists(train_y_path) and os.path.exists(test_x_path) and os.path.exists(test_y_path):
        print('--------------------Load Datasets--------------------')
        train_x_data = np.load(train_x_path)  # 加载训练数据
        train_x_data = np.reshape(train_x_data, (len(train_x_data), 32, 32))  # 将数据调整为32x32的大小
        train_y_data = np.load(train_y_path)  # 加载训练标签

        test_x_data = np.load(test_x_path)  # 加载测试数据
        test_x_data = np.reshape(test_x_data, (len(test_x_data), 32, 32))  # 调整为32x32大小
        test_y_data = np.load(test_y_path)  # 加载测试标签
    else:  # 如果数据不存在，生成数据
        print('--------------------Generate Datasets--------------------')
        train_x_data, train_y_data = generateds(data_train_path)  # 生成训练数据
        test_x_data, test_y_data = generateds(data_test_path)  # 生成测试数据

        print('--------------------Save Datasets--------------------')
        # 随机打乱训练数据
        np.random.seed(116)
        np.random.shuffle(train_x_data)
        np.random.seed(116)
        np.random.shuffle(train_y_data)

        train_x_data = np.reshape(train_x_data, (len(train_x_data), 32, 32))  # 调整为32x32大小
        np.save(train_x_path, train_x_data)  # 保存训练数据
        np.save(train_y_path, train_y_data)  # 保存训练标签

        # 随机打乱测试数据
        np.random.seed(116)
        np.random.shuffle(test_x_data)
        np.random.seed(116)
        np.random.shuffle(test_y_data)

        test_x_data = np.reshape(test_x_data, (len(test_x_data), 32, 32))  # 调整为32x32大小
        np.save(test_x_path, test_x_data)  # 保存测试数据
        np.save(test_y_path, test_y_data)  # 保存测试标签

    # 将加载的 NumPy 数组赋给训练和测试数据
    X_train, X_test = train_x_data, test_x_data
    Y_train, Y_test = train_y_data, test_y_data

    # 设置DataLoader的额外参数
    kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}

    # 定义训练和测试数据加载器
    train_dataset = MyDataset(X_train, Y_train)
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, **kwargs)

    test_dataset = MyDataset(X_test, Y_test)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, **kwargs)

    # 创建神经网络模型
    model = Net()
    if args.cuda:
        model.cuda()  # 如果使用CUDA，模型移至GPU

    # 定义优化器
    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

    # 开始训练和测试
    for epoch in range(1, args.epochs + 1):
        train(epoch, model, train_loader, optimizer)  # 调用训练函数
    torch.save(model.state_dict(), 'save.pt')  # 将模型的状态字典（参数）保存到文件
    print("Model saved to 'save.pt'")
    test(model, test_loader)  # 调用测试函数
